from tools.camera.search_camera import setCamera
from multiprocessing import Queue
import threading
import cv2
from tools.log import log
from ai_library.components.config import cfg as ai_cfg
from tools.databus.bus import workFlow

class ipcThread(threading.Thread):
    def __init__(self, camera="rtsp", state_dict:dict=None,
                 q_src:Queue=None, q_infer:Queue=None):
        threading.Thread.__init__(self)
        self.cap_ipc = setCamera(camera)  # 网络摄像头
        self.state_dict = state_dict      # 全局状态dict获取识别结果
        self.q_src = q_src                # 消息队列传递原始图到识别插件
        self.q_infer = q_infer            # 消息队列传递绘制识别结果后的图像到web显示插件

    def run(self):
        while True:
            if self.cap_ipc != "":
                ret, ipc_frame = self.cap_ipc.read()
                ipc_frame = cv2.resize(ipc_frame, (ai_cfg["cam_width"], ai_cfg["cam_height"]))

                cv2.imshow("ipc_frame", ipc_frame)
                if not self.q_src.full() and not ipc_frame is None:
                    self.q_src.put(bytearray(ipc_frame))

                # 识别结果绘制


                if not self.q_infer.full() and not ipc_frame is None:
                    self.q_infer.put(bytearray(ipc_frame))
                cv2.waitKey(10)


class csiThread(threading.Thread):
    def __init__(self, camera="mipi", state_dict:dict=None,
                 q_src:Queue=None, q_infer:Queue=None):
        threading.Thread.__init__(self)
        self.cap_csi = setCamera(camera)  # 网络摄像头
        self.state_dict = state_dict      # 全局状态dict获取识别结果
        self.q_src = q_src                # 消息队列传递原始图到识别插件
        self.q_infer = q_infer            # 消息队列传递绘制识别结果后的图像到web显示插件

    def run(self):
        while True:
            if self.cap_csi != "":
                ret, ipc_frame = self.cap_csi.read()
                ipc_frame = cv2.resize(ipc_frame, (ai_cfg["cam_width"], ai_cfg["cam_height"]))

                if not self.q_src.full() and not ipc_frame is None:
                    self.q_src.put(bytearray(ipc_frame))

                # 识别结果绘制


                if not self.q_infer.full() and not ipc_frame is None:
                    self.q_infer.put(bytearray(ipc_frame))

class ImageGet():
    def __init__(self):
        pass

    def onExit(self):
        pass

    def worker(self, state_dict:dict=None,
                 q_ipc_src:Queue=None,
                 q_ipc_infer:Queue=None,
                 q_csi_src:Queue=None,
                 q_csi_infer:Queue=None):

        self.ipc_thread = ipcThread("0", state_dict, q_ipc_src, q_ipc_infer)
        self.csi_thread = csiThread("mipi", state_dict, q_csi_src, q_csi_infer)
        log.info("图像获取线程已启动！！")
        self.ipc_thread.start()
        self.csi_thread.start()


def imgGetPluginRegist(state_dict:dict=None,
                 q_ipc_src:Queue=None,
                 q_ipc_infer:Queue=None,
                 q_csi_src:Queue=None,
                 q_csi_infer:Queue=None):

    img_get_plugin = ImageGet()

    workFlow.registBus(img_get_plugin.__class__.__name__,
                       img_get_plugin.worker,
                       (state_dict, q_ipc_src, q_ipc_infer, q_csi_src, q_csi_infer),
                       img_get_plugin.onExit)


if __name__ == "__main__":
    import multiprocessing as mp

    q_ipc = mp.Queue(1)
    q_ipc_if = mp.Queue(1)
    imageGet = ImageGet()
    imageGet.run(q_ipc_src=q_ipc, q_ipc_infer=q_ipc_if)